import time
import argparse
import torch
from w2v_costh import backward as bkwd


def parse_args():
    parser = argparse.ArgumentParser(description="hyper-params of `w2v_costh` module")
    parser.add_argument("--batch-size", type=int, default=3, help="batch size of input tensor")
    parser.add_argument("--weight-dim0", type=int, default=768, help="the first dimension of weight")
    parser.add_argument("--weight-dim1", type=int, default=1211, help="the second dimension of weight")
    args = parser.parse_args()

    return args


if __name__ == "__main__":
    args = parse_args()
    assert args.weight_dim0 == 768, "the first dimension of weight must be equal to 768"
    assert args.weight_dim1 == 1211, "the second dimension of weight must be equal to 1211"

    bkwd_total_time = 0.0
    bkwd_grad_mm = torch.randn(args.batch_size, args.weight_dim1).cuda(0)
    bkwd_x_norm = torch.randn(args.batch_size, args.weight_dim0).cuda(0)
    bkwd_w = torch.randn(args.weight_dim0, args.weight_dim1).cuda(0)
    bkwd_costh_out = torch.randn(args.weight_dim0, args.weight_dim1).cuda(0)
    bkwd_norm_vals = torch.randn((args.weight_dim1, )).cuda(0)

    for idx in range(11000):
        cond = idx >= 1000
        if cond:
            bkwd_start_time = time.time()
        bkwd_res = bkwd(bkwd_grad_mm, bkwd_x_norm, bkwd_w, bkwd_costh_out, bkwd_norm_vals)
        if cond:
            bkwd_total_time += time.time() - bkwd_start_time
    print(f"Backward time of inner function is {bkwd_total_time * 100:.5f} us / iter.")

